package com.example.demo.controller;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import com.example.demo.domain.vo.ChatVO;

import jakarta.annotation.Resource;
import reactor.core.publisher.Flux;

@RequestMapping("/chat")
@RestController
public class ChatController {
    Logger logger = LoggerFactory.getLogger(ChatController.class);

    private static final Map<String, Boolean> GENERATE_STATUS = new ConcurrentHashMap<>();

    @Resource
    private ChatClient chatClient;

    @Resource
    private VectorStore vectorStore;

    /**
     * ChatClient 流式调用
     */
    @GetMapping(value = "/wt", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ChatVO> streamChat(@RequestParam("question") String question,
            @RequestParam("sessionId") String sessionId) {
        var qaAdvisor = QuestionAnswerAdvisor.builder(vectorStore)
                .searchRequest(SearchRequest.builder().similarityThreshold(0.75d).topK(3).build())
                .build();

        return chatClient.prompt()
                .advisors(advisor -> advisor
                        .advisors(qaAdvisor)
                        .param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, sessionId))
                .user(question)
                .stream()
                .chatResponse()
                .doFirst(() -> {
                    GENERATE_STATUS.put(sessionId, true);
                    logger.info("开始流式调用");
                })
                .doOnError(throwable -> GENERATE_STATUS.remove(sessionId))
                .takeWhile(s -> GENERATE_STATUS.getOrDefault(sessionId, false))
                .doFinally(signalType -> {
                    GENERATE_STATUS.put(question, false);
                    logger.info("流式调用结束");
                })
                .map(data -> {
                    return new ChatVO().builder()
                            .result(data.getResult().getOutput().getText())
                            .eventType("question")
                            .build();
                })
                .concatWith(Flux.just(
                        new ChatVO().builder()
                                .result("[STOP]")
                                .eventType("end")
                                .build()));
    }

    @PostMapping("/stop")
    public void stop(@RequestParam("sessionId") String sessionId) {
        logger.info("停止流式调用");
        GENERATE_STATUS.remove(sessionId);
    }

    /**
     * 搜索向量数据库
     *
     * @param query 搜索关键字
     */
    @GetMapping("/search")
    public List<Document> search() {
        logger.info("开始搜索向量数据库");
        return this.vectorStore.similaritySearch(SearchRequest.builder()
                .query("芒果") // 设置查询条件
                .topK(1) // 设置最多返回的文档数量
                .build());
    }
}
